extern "C"
{
#include "types.h"
#include "s0.h"
#include "field.h"
#include "rotate.h"
#include "h_func.h"
#include "expandkey.h"
#include "error_insert_help.h"
#include "set_error_insert.h"
#include "twofish.h"


byte gf_256_169_mul(byte,byte);
}


#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <ctime>
#include <vector>
#include <unordered_set>
#include <map>

using namespace std;


//output_differ adder_offset input_differ
byte add_in_diff[0x100][0x100][0x100];
int add_in_diff_size[0x100][0x100];

const char* file_name = "od_ad_id.rlt";
void save_to_file()
{
	FILE* save_file = fopen(file_name, "wb+");
	if(save_file == NULL)
	{
		printf("can't open file %s\n", file_name);
		exit(1);
	}

	int try_od,try_ad;
	for(try_od = 0; try_od <= 0xff; ++try_od)
	{
		for(try_ad = 0; try_ad <= 0xff; ++try_ad)
		{
			byte out_diff = try_od;
			int try_in_diff,try_adder;
			add_in_diff_size[try_od][try_ad] = 0;
			for(try_in_diff = 0; try_in_diff <= 0xff; ++try_in_diff)
			{
				int found = 0;
				int try_in;
				for(try_in = 0; try_in <= 0xff; ++try_in)
				{
					int try_in_e = try_in ^ try_in_diff;
					for(try_adder = 0; try_adder <= 0xff; ++try_adder)
					{
						int try_adder_e = try_adder + try_ad;
						byte out = try_in + try_adder;
						byte out_e = try_in_e + try_adder_e;
						byte od = out ^ out_e;
						if(od == out_diff)
						{
							found = 1;
							break;
						}
					}
					if(found)
						break;
				}
				if(found)
				{
					int size = add_in_diff_size[try_od][try_ad];
					add_in_diff[try_od][try_ad][size] = try_in_diff;
					++add_in_diff_size[try_od][try_ad];
				}
			}
		}
		printf("try_od:%d\n", try_od);
	}

	for(try_od = 0; try_od <= 0xff; ++try_od)
	{
		for(try_ad = 0; try_ad <= 0xff; ++try_ad)
		{
			int size = add_in_diff_size[try_od][try_ad];
			int write_size = fwrite(&size, sizeof(size), 1, save_file);
			if(write_size != 1)
			{
				printf("fwrite error!\n");
				exit(1);
			}
		}
	}

	for(try_od = 0; try_od <= 0xff; ++try_od)
	{
		for(try_ad = 0; try_ad <= 0xff; ++try_ad)
		{
			int size = add_in_diff_size[try_od][try_ad];
			int j;
			for( j = 0; j < size; ++j)
			{
				byte value = add_in_diff[try_od][try_ad][j];
				int write_size = fwrite(&value, sizeof(value), 1, save_file);
				if(write_size != 1)
				{
					printf("fwrite error!\n");
					exit(1);
				}
			}
		}
	}

	fclose(save_file);
}

void read_from_file()
{
	FILE* save_file = fopen(file_name, "rb");
	if(save_file == NULL)
	{
		printf("open file %s to read error\n", file_name);
		exit(1);
	}

	int try_od, try_ad;
	for(try_od = 0; try_od <= 0xff; ++try_od)
	{
		for(try_ad = 0; try_ad <= 0xff; ++try_ad)
		{
			int read_count = fread(&add_in_diff_size[try_od][try_ad], sizeof(int), 1, save_file);
			if(read_count != 1)
			{
				printf("fread error!\n");
				exit(1);
			}
		}
	}

	for(try_od = 0; try_od <= 0xff; ++try_od)
	{
		for(try_ad = 0; try_ad <= 0xff; ++try_ad)
		{
			int size = add_in_diff_size[try_od][try_ad];
			int j;
			for( j = 0; j < size; ++j)
			{
				int read_count = fread(&add_in_diff[try_od][try_ad][j], sizeof(byte), 1, save_file);
				if(read_count != 1)
				{
					printf("fread error!\n");
					exit(1);
				}
			}
		}
	}

	fclose(save_file);
}

bool test_pass(int pos, byte in_diff, byte out_diff, vector<vector<unordered_set<byte> > >& aid)
{
	static bool pos_initialled = false;
	static vector<vector<int> > test_pos;
	if(!pos_initialled)
	{
		static int ca_pos[6] = {
			1,2,3,5,6,7
		};
		pos_initialled = true;
		test_pos.resize(8);
		test_pos[0].push_back(0);
		test_pos[4].push_back(0);

		for(int i = 0; i < sizeof(ca_pos)/sizeof(ca_pos[0]); ++i)
		{
			test_pos[ca_pos[i]].push_back(0);
			test_pos[ca_pos[i]].push_back(1);
//			test_pos[ca_pos[i]].push_back(2);
		}
	}
	for(int i = 0; i < test_pos[pos].size(); ++i)
	{
		unordered_set<byte>& test_ad = aid[test_pos[pos][i]][out_diff];
		if(test_ad.find(in_diff) != test_ad.end())
			return true;
	}
	return false;
}


class print_runtime
{
public:
	print_runtime()
	{
		start_time = clock();
	}
	~print_runtime()
	{
		clock_t end_time = clock();
		printf("Cost time:%f\n", (end_time-start_time)/(double)CLOCKS_PER_SEC);
	}
private:
	clock_t start_time;
};

struct cand_struct
{
	cand_struct(byte o, byte* s, int k_s):
		owk(o), key_size(k_s)
	{
		memcpy(sk, s, key_size);
	}
	cand_struct():
		owk(0),key_size(0)
	{
	}
	byte owk;
	byte sk[4];
	int key_size;
};

#include <stdexcept>

class cand_generator
{
public:
	cand_generator(vector<cand_struct>& cand,int key_size):
		m_cand(cand),m_key_size(key_size),m_owk_tra(0),m_sk_tra(0),m_it(m_cand.begin())
	{
		m_sk_max = ((long long)1) << (8*m_key_size);
	}
	cand_struct operator()()
	{
		cand_struct result;
		if(!has_next())
		{
			throw runtime_error("cand generator does not has next to generate!\n");
		}
		if(m_cand.size() == 0)
		{
			result.owk = m_owk_tra;
			memcpy(result.sk, &m_sk_tra, 4);
/*			for(int i = 0; i < m_key_size; ++i)
			{
				result.sk[i] = (m_sk_tra >> (8*i)) & 0xff;
			}*/
			result.key_size = m_key_size;
			m_sk_tra++;
			if(m_sk_tra == m_sk_max)
			{
				m_sk_tra = 0;
				m_owk_tra++;
//				printf("owk:%x\n",m_owk_tra);
			}
//			if(m_sk_tra % 0x10000000 == 0)
//				printf("sk:%x\n", m_sk_tra);
		}
		else
		{
			result = *m_it;
			m_it++;
		}
		return result;
	}
	bool has_next() const
	{
		if(m_cand.size() == 0)
			return (m_owk_tra < 0x100);
		else
			return m_it != m_cand.end();
	}
private:
	vector<cand_struct> &m_cand;
	int m_key_size;
	int m_owk_tra;
	long long m_sk_tra;
	long long m_sk_max;
	vector<cand_struct>::iterator m_it;
};

int main()
{
	int key_size;
	printf("Please input the attack key size:");
	int ret = scanf("%d", &key_size);
	if(key_size != 16 &&
		key_size != 24 && 
		key_size != 32)
	{
		printf("The key size must be 16,24 or 32!\n");
		return 1;
	}

	byte plaintext[BLOCK_SIZE] = {
		0
	};
	byte cipher[BLOCK_SIZE];
	byte error_cipher[BLOCK_SIZE];

	key_type encrypt_key;
	byte key[32] = {
		0
	};
	pre_expand_key(key, key_size, &encrypt_key);

	read_from_file();

	vector<cand_struct> cand;
	cand.reserve(0xffffff);
/*	for(int i = 0; i <= 0xffffff; ++i)
	{
		cand_t.push_back(i);
	}*/

	vector<vector<unordered_set<byte> > > aid;
	aid.resize(0x100);
	for(int i = 0; i < aid.size(); ++i)
	{
		aid[i].resize(0x100);
		for(int out_diff = 0; out_diff <= 0xff; ++out_diff)
		{
			for(int in_j = 0; in_j < add_in_diff_size[out_diff][i]; ++in_j)
			{
				aid[i][out_diff].insert(add_in_diff[out_diff][i][in_j]);
			}
		}
	}

	srand(time(NULL));

	int try_times = 100;
	map<int, vector<int> > sc2cs;

	vector<vector<int> > equ_cs;

	printf("Sample count\t\tcand size\n");
	for(int try_count = 0; try_count < try_times; ++try_count)
	{

///////////////////////////////////////////////////////////////////////////
///////////////////////////////////////////////////////////////////////////
///////////////////////////////////////////////////////////////////////////
///////////////////////////////////////////////////////////////////////////
///////////////////////////////////////////////////////////////////////////
///////////////////////////////////////////////////////////////////////////
	for(int i = 0; i < key_size; ++i)
	{
		key[i] = rand() & 0xff;
	}

	printf("Master key:\n");
	for(int i = 0; i < key_size; ++i)
	{
		printf("%x", key[i]);
	}
	printf("\n");

	pre_expand_key(key, key_size, &encrypt_key);
	cand.resize(0);
	
	int sample_count = 0;
//	print_runtime prt;
	while(true)
	{
		sample_count++;

		clear_error_struct();

		plaintext[0] = rand();
		byte xor_error;
		xor_error = rand() % 255 + 1;
		error_struct error_data;
		error_data.size = sizeof(error_data);
		error_data.byte_pos = 0;
		error_data.round = 15;
		error_data.error_byte = xor_error;
		error_data.apply_func = xor_error_apply;

//		printf("error byte:%.2x\n", xor_error);

		if(!add_error_struct(&error_data))
		{
			printf("add_error_struct error!\n");
			return 1;
		}
		
		int i;

		set_error_insert(NULL);
		twofish_128_encrypt(plaintext, cipher, &encrypt_key);

		set_error_insert(error_insert_hook_func);
		twofish_128_encrypt(plaintext, error_cipher, &encrypt_key);

		byte out_diff[8];
		for(int i = 0; i < 8; ++i)
		{
			out_diff[i] = cipher[8+i] ^ error_cipher[8+i];
		}
		word tmp_word = join_word(out_diff);
		tmp_word = rol(tmp_word, 1);
		split_word(tmp_word, out_diff, 4);
//		printf("od:%.2x\n", out_diff[0]);

		byte rt0 = cipher[0];
		byte rt0_e = error_cipher[0];

//		printf("st0:%.2x,st0_e:%.2x\n", rt0, rt0_e);
		vector<cand_struct> tmp_result;
		tmp_result.reserve(cand.size());
//		for(int i = 0; i < cand.size(); ++i)
		cand_generator gen(cand,key_size/8);
		while(gen.has_next())
		{
			cand_struct cand_s = gen();
/*			int val = cand[i];
			byte sk0 = val & 0xff;
			byte sk1 = (val >> 8) & 0xff;
			byte sk[] ={
				sk0,sk1
			};
			byte owk0 = (val >> 16) & 0xff;*/
			byte owk0 = cand_s.owk;
			
			rt0  = cipher[0] ^ owk0;
			rt0_e = error_cipher[0] ^ owk0;

			rt0 = s0(rt0, cand_s.sk, cand_s.key_size);
			rt0_e = s0(rt0_e, cand_s.sk, cand_s.key_size);
			
			static const byte coeff[8] = {
				0x01, 0x5B, 0xEF, 0xEF,
				0x01, 0x5B, 0xEF, 0xEF,
			};

			bool pass = true;
			byte rt,rt_e;
			for(int j = 0; j < 8; ++j)
			{
				rt = gf_256_169_mul(coeff[j], rt0);
				rt_e = gf_256_169_mul(coeff[j], rt0_e);

				byte in_diff = rt ^ rt_e;
				if(!test_pass(j, in_diff, out_diff[j], aid))
				{
					pass = false;
					break;
				}
			}
			if(pass)
			{
				tmp_result.push_back(cand_s);
			}
		}

		cand.swap(tmp_result);


		if(sc2cs.find(sample_count) == sc2cs.end())
		{
			sc2cs[sample_count] = vector<int>();
		}
		sc2cs[sample_count].push_back(cand.size());

//		printf("Sample count:%d,cand size:%d\n", sample_count, cand.size());
		printf("%d\t\t\t%d\n", sample_count, cand.size());	
		if(cand.size() == 1)
		{
			cand_struct& cand_s = cand[0];
			byte owk0 = cand_s.owk;
			printf("Found!owk0:%.2x,", owk0);
			for(int i = 0; i < cand_s.key_size; ++i)
			{
				printf("sk%d:%.2x,", i, cand_s.sk[i]);
			}
			printf("\n");
			break;
		}
		else if(cand.size() == 0)
		{
			printf("Cand size is zero!error!!!!!!!!!!!!!!!!!!\n");
			return 1;
		}
	}
///////////////////////////////////////////////////////////////////////////
///////////////////////////////////////////////////////////////////////////
///////////////////////////////////////////////////////////////////////////
///////////////////////////////////////////////////////////////////////////
///////////////////////////////////////////////////////////////////////////
	}

	typedef map<int, vector<int> >::iterator it_type;
	for(it_type it = sc2cs.begin(); it != sc2cs.end(); ++it)
	{
		vector<int>& cs = it->second;
		float ave = cs[0];
		for(int i = 1; i < cs.size(); ++i)
		{
			ave = ave * i/(i+1) + cs[i] / (i+1);
		}
		printf("sample count:%d,average is:%d\n", it->first, (int)ceil(ave));
	}

	return 0;
}
